#include "Lib.h"
#include "Cache.h"
#include <unordered_map>
#include <unordered_set>

using namespace std;

typedef unordered_map<jmethodID, MethodInfo*> MethodCacheType;
MethodCacheType MethodCache;

MethodInfo* GetCachedMethodInfo(jmethodID method)
{
	auto it = MethodCache.find(method);
	if (it == MethodCache.end())
		return NULL;
	return it->second;
}

bool TryAddCacheMethodInfo(jmethodID method, InterestMethodTableEntry* list)
{
	if (MethodCache.find(method) != MethodCache.end())
		return true;

	MethodInfo info;
	info.MethodID = method;

	// get method information
	jvmtiError error = GJvmti->GetMethodName(method, (char**)&info.Name,
		(char**)&info.Signature, NULL);
	CheckError(error, "GetMethodName");

	// get class information
	jclass methodClass = NULL;

	error = GJvmti->GetMethodDeclaringClass(method, &methodClass);
	CheckError(error, "GetMethodDeclaringClass");
	error = GJvmti->GetClassSignature(methodClass,
		(char**)&info.ClassSignature, NULL);
	CheckError(error, "GetClassSignature");

	// find existence in list
	int i = 0;
	for (; list[i].Name != NULL; i++)
	{
		if (strcmp(list[i].Name, info.Name) != 0)
			continue;
		if (list[i].Signature &&
			strcmp(list[i].Signature, info.Signature) != 0)
			continue;
		if (list[i].ClassSignature &&
			strcmp(list[i].ClassSignature, info.ClassSignature) != 0)
			continue;
		break;
	}

	// not found, free memory
	if (list[i].Name == NULL)
	{
		GJvmti->Deallocate((unsigned char*)info.Name);
		GJvmti->Deallocate((unsigned char*)info.Signature);
		GJvmti->Deallocate((unsigned char*)info.ClassSignature);
		return false;
	}

	info.CallbackEntry = list[i].CallbackEntry;
	info.CallbackExit = list[i].CallbackExit;

	MethodCache.insert(make_pair(method, new MethodInfo(info)));
	return true;
}

void AddCacheMethodInfo(jmethodID method,
	const char* name, const char* sig, const char* classsig, callbackMethod callback)
{
	MethodInfo* minfo = new MethodInfo;
	minfo->MethodID = method;
	minfo->Name = name;
	minfo->Signature = sig;
	minfo->ClassSignature = classsig;
	minfo->CallbackEntry = callback;
	minfo->CallbackExit = callback;

	MethodCache.insert(make_pair(method, minfo));
}

bool TryAddCacheMethodInfo(jmethodID method,
	const char* name, const char* sig, const char* classsig,
	InterestMethodTableEntry* list)
{
	MethodInfo info;
	info.MethodID = method;
	info.Name = name;
	info.Signature = sig;
	info.ClassSignature = classsig;

	// find existence in list
	int i = -1;
	while (list[i + 1].Name != NULL)
	{
		i++;
		if (strcmp(list[i].Name, info.Name) != 0)
			continue;
		if (list[i].Signature && strcmp(list[i].Signature, info.Signature) != 0)
			continue;
		if (list[i].ClassSignature && strcmp(list[i].ClassSignature, info.ClassSignature) != 0)
			continue;
	}

	// not found, free memory
	if (list[i].Name == NULL)
		return false;

	info.CallbackEntry = list[i].CallbackEntry;
	info.CallbackExit = list[i].CallbackExit;

	MethodCache.insert(make_pair(method, new MethodInfo(info)));

	return true;
}

// classSignature cache
typedef unordered_set<jint> ClassCacheType;
ClassCacheType ClassCache;

bool CheckCachedClassInfo(jclass klass)
{
	jint classHash = 0;
	jvmtiError error = GJvmti->GetObjectHashCode(klass, &classHash);
	assert(error == JVMTI_ERROR_NONE);

	if (ClassCache.find(classHash) != ClassCache.end())
		return true;

	// get signature
	char* classSignature = NULL;
	error = GJvmti->GetClassSignature(klass, &classSignature, NULL);
	CheckError(error, "cannot get class signature");

	// get signature
	char* classSource = NULL;
	error = GJvmti->GetSourceFileName(klass, &classSource);
	CheckError(error, "cannot get class sourcefile");

	// dump to file
	LogClasses->BeginRecord("%08X,%s,%s", classHash, classSignature,
		classSource == NULL ? "" : classSource);
	LogClasses->EndRecord();

	GJvmti->Deallocate((unsigned char*)classSignature);
	GJvmti->Deallocate((unsigned char*)classSource);

	ClassCache.insert(classHash);
	return false;
}

// method SourceMap table
typedef unordered_set<jmethodID> MSourceCacheType;
MSourceCacheType MSourceCache;
bool CheckCachedMethodSourceInfo(jmethodID method)
{
	if (MSourceCache.find(method) != MSourceCache.end())
		return true;

	// <methodID>,<methodName>,<methodSignature>,<declaringClassHash>,[lineTable]
	// get method information
	char* methodName = NULL;
	char* methodSignature = NULL;
	jvmtiError error = GJvmti->GetMethodName(method, &methodName,
		&methodSignature, NULL);
	CheckError(error, "GetMethodName");

	// get class information
	jclass methodClass = NULL;
	jint classHash = 0;

	error = GJvmti->GetMethodDeclaringClass(method, &methodClass);
	CheckError(error, "GetMethodDeclaringClass");
	error = GJvmti->GetObjectHashCode(methodClass, &classHash);
	CheckError(error, "GetMethodDeclaringClassHashCode");

	// dump
	LogMethods->BeginRecord("%08X,%s,%s,%08X", method,
		methodName, methodSignature, classHash);

	// get line tables
	jint lineCount = -1;
	jvmtiLineNumberEntry* lineTable = NULL;
	error = GJvmti->GetLineNumberTable(method, &lineCount,
		&lineTable);
	if (error == JVMTI_ERROR_NONE)
	{
		LogMethods->Record(",%d", lineCount);
		for (int i = 0; i < lineCount; i++)
			LogMethods->Record(";%08X;%08X",
			lineTable[i].line_number,
			lineTable[i].start_location);
		GJvmti->Deallocate((unsigned char*)lineTable);
	}
	else
		LogMethods->Record(",0"); // no line table

	LogMethods->EndRecord();

	GJvmti->Deallocate((unsigned char*)methodName);
	GJvmti->Deallocate((unsigned char*)methodSignature);

	// check class cache
	CheckCachedClassInfo(methodClass);

	MSourceCache.insert(method);

	return false;
}